#单眼
import torch
from torchvision import transforms
from PIL import Image
import numpy as np
from model.model import MultiLabelBiFPNResNet
# 定义类别标签（顺序与模型输出一致）
labels = ['N', 'D', 'G', 'C', 'A', 'H', 'M', 'O']
# 定义标签映射
label_mapping = {
    'N': '正常',
    'D': '糖尿病',
    'G': '青光眼',
    'C': '白内障',
    'A': 'AMD',
    'H': '高血压',
    'M': '近视',
    'O': '其他疾病/异常'
}

# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 根据训练时输入尺寸调整
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# 加载模型（全局变量，避免重复加载）
model = None

def load_model():
    """加载预训练模型"""
    global model
    if model is None:
        try:
            print("初始化模型...")
            model = MultiLabelBiFPNResNet(num_classes=8)
            
            # 使用绝对路径加载模型
            import os
            model_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "best_resnet_model.pth")
            print(f"加载模型文件: {model_path}")
            
            # 加载模型权重
            state_dict = torch.load(model_path, map_location=torch.device('cpu'))
            # 打印模型结构信息
            print("\n当前模型结构中的键:")
            current_keys = set(model.state_dict().keys())
            print(sorted(current_keys))
            
            print("\n加载文件中的键:")
            loaded_keys = set(state_dict.keys())
            print(sorted(loaded_keys))

            # 创建新的状态字典，处理前缀差异
            new_state_dict = {}
            for k, v in state_dict.items():
                # 将 'resnet' 前缀替换为 'backbone'
                if k.startswith('resnet.'):
                    new_key = k.replace('resnet.', 'backbone.')
                    new_state_dict[new_key] = v
                else:
                    new_state_dict[k] = v
            
            # 检查键的匹配情况
            current_keys = set(model.state_dict().keys())
            new_keys = set(new_state_dict.keys())
            missing_keys = current_keys - new_keys
            unexpected_keys = new_keys - current_keys
            
            if missing_keys:
                print("\n缺失的键:")
                print(sorted(missing_keys))
            if unexpected_keys:
                print("\n多余的键:")
                print(sorted(unexpected_keys))
            
            # 加载模型
            try:
                model.load_state_dict(new_state_dict, strict=True)
                print("\n模型加载成功（严格模式）")
            except Exception as e:
                print(f"\n严格加载失败: {str(e)}")
                print("尝试非严格模式加载...")
                model.load_state_dict(new_state_dict, strict=False)
                print("模型加载成功（非严格模式）")
            
            model.eval()
            print("模型设置为评估模式")
            return model
            
        except Exception as e:
            import traceback
            print(f"\n模型加载过程中发生错误: {str(e)}")
            print(traceback.format_exc())
            raise
    return model

def predict_single_eye(image_path, threshold=0.5):
    """
    对单眼图像进行疾病预测
    
    参数:
        image_path (str): 眼底图像的路径
        threshold (float): 预测阈值，默认为0.5
        
    返回:
        dict: 包含预测结果的字典，包括各类别概率、预测的疾病和最终诊断
    """
    try:
        print(f"开始单眼预测: {image_path}")
        
        # 加载模型
        model = load_model()
        
        # 加载并预处理图像
        image = Image.open(image_path).convert("RGB")
        image_tensor = transform(image).unsqueeze(0)  # 增加 batch 维度
        
        # 预测
        with torch.no_grad():
            # 对于单眼预测，我们将同一图像作为左右眼输入
            outputs = model(image_tensor, image_tensor)
            probabilities = torch.sigmoid(outputs).squeeze().cpu().numpy()
        
        # 根据阈值确定预测结果
        predicted = (probabilities >= threshold).astype(int)
        
        # 获取预测的疾病
        predicted_diseases = []
        for i, pred in enumerate(predicted):
            if i > 0 and pred == 1:  # 跳过正常标签(N)，只关注疾病
                predicted_diseases.append(labels[i])
        
        # 确定最终诊断
        if len(predicted_diseases) == 0:
            diagnosis = label_mapping['N']
        else:
            diagnosis = "、".join([label_mapping[d] for d in predicted_diseases])
        
        print(f"预测完成，诊断结果: {diagnosis}")
        
        # 构建结果字典
        result = {
            "probabilities": {label_mapping[labels[i]]: float(probabilities[i]) for i in range(len(labels))},
            "predicted_diseases": [label_mapping[d] for d in predicted_diseases],
            "diagnosis": diagnosis
        }
        
        return result
    
    except Exception as e:
        import traceback
        print(f"单眼预测失败: {str(e)}")
        print(traceback.format_exc())
        return {"error": str(e)}

# 如果直接运行此脚本，则执行示例预测
if __name__ == "__main__":
    # 示例用法
    image_path = "path/to/your/image.jpg"  # 替换为实际图像路径
    result = predict_single_eye(image_path)
    
    print("各类别概率:")
    for disease, prob in result["probabilities"].items():
        print(f"类别 {disease}: 概率 {prob:.4f}")
    
    print("\n最终诊断结果:")
    print(result["diagnosis"])




